import numpy as np
from common.geoimage.raster_dataset import RasterDataset
from common.geoimage.raster_dataset import SceneMeta
from common.logger import logger


def extract_time_data(
    band_idx,
    src_filepath,
    dest_filepath,
    extract_mode="single",
    start_band=0,
    end_band=-1,
    override=False,
):
    from scipy.stats import mode

    if extract_mode == "single":
        assert len(band_idx) == 1, "Please only provide the sinlge band index"
        src_rst = RasterDataset.from_file(src_filepath, bands=band_idx)
        if not os.path.exists(dest_filepath) or override:
            src_rst.to_geotiff(dest_filepath)
        else:
            logger.info(f"{dest_filepath} already existed. Will skip")
    else:
        dest_filepath = dest_filepath.replace(
            "_".join([str(val) for val in band_idx]), "majority"
        )
        os.makedirs(os.path.dirname(dest_filepath), exist_ok=True)
        src_rst = RasterDataset.from_file(src_filepath)
        src_rst.data = src_rst.data[start_band:end_band]
        num_class = [0] + list(np.unique(src_rst.data))
        prob_avg = np.zeros(
            [len(num_class), src_rst.data.shape[1], src_rst.data.shape[2]]
        )
        for idx, c in enumerate(num_class):
            prob_avg[idx] = (src_rst.data == c).mean(axis=0)
        src_rst.data = np.argmax(prob_avg, axis=0).astype(np.uint16)
        src_rst.meta = src_rst.meta.update(n_bands=1)
        if not os.path.exists(dest_filepath) or override:
            src_rst.to_geotiff(dest_filepath)
        else:
            logger.info(f"{dest_filepath} already existed. Will skip")


if __name__ == "__main__":
    import os
    import json
    from multiprocessing.pool import ThreadPool
    import analytics.crop_recognition.common_settings.model_config as config

    src_folder = (
        f"/your_path_to_load_inference/{config.region}/{config.name}"  # noqa:E501
    )
    dest_folder = f"/your_path_to_load_inference/{config.region}/{config.name}/processed"  # noqa:E501

    band_idx = [10]
    year_list = [2023]

    dest_folder_name = "_".join([str(val) for val in band_idx])
    dest_folder_path = os.path.join(dest_folder, dest_folder_name)
    if not os.path.exists(dest_folder_path):
        os.makedirs(dest_folder_path)

    for year in year_list:
        year_str = str(year)
        year_int = int(year)
        cur_src_year_folder = os.path.join(src_folder, year_str)
        cur_dest_year_folder = os.path.join(dest_folder_path, year_str)
        if not os.path.exists(cur_dest_year_folder):
            os.makedirs(cur_dest_year_folder)
        pool = ThreadPool(30)
        for root, _, filenames in os.walk(cur_src_year_folder):
            for filename in filenames:
                if filename.endswith(".json"):
                    cur_json_filepath = os.path.join(root, filename)
                    with open(cur_json_filepath, "r") as f:
                        cur_json_dict = json.load(f)

                    cur_class_tif_name = cur_json_dict["image_assets"]["class"]
                    cur_confidence_name = cur_json_dict["image_assets"]["confidence"]
                    cur_class_tif_path = os.path.join(root, cur_class_tif_name)
                    dest_class_tif_path = os.path.join(
                        dest_folder_path,
                        os.path.relpath(cur_class_tif_path, src_folder),
                    )

                    cur_conf_tif_path = os.path.join(root, cur_confidence_name)
                    dest_conf_tif_path = os.path.join(
                        dest_folder_path, os.path.relpath(cur_conf_tif_path, src_folder)
                    )

                    dest_json_path = os.path.join(
                        dest_folder_path, os.path.relpath(cur_json_filepath, src_folder)
                    )
                    with open(dest_json_path, "w") as f:
                        json.dump(cur_json_dict, f, indent=4)

                    pool.apply_async(
                        extract_time_data,
                        args=[
                            band_idx,
                            cur_class_tif_path,
                            dest_class_tif_path,
                        ],
                    )

        pool.close()
        pool.join()
